// Copyright (c) 2016 GitHub, Inc.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

#include "atom/browser/relauncher.h"

#include <windows.h>

#include "base/process/launch.h"
#include "base/strings/stringprintf.h"
#include "base/win/scoped_handle.h"
#include "sandbox/win/src/nt_internals.h"
#include "sandbox/win/src/win_utils.h"
#include "ui/base/win/shell.h"

namespace relauncher {

namespace internal {

    namespace {

        const CharType* kWaitEventName = L"ElectronRelauncherWaitEvent";

        HANDLE GetParentProcessHandle(base::ProcessHandle handle)
        {
            NtQueryInformationProcessFunction NtQueryInformationProcess = nullptr;
            ResolveNTFunctionPtr("NtQueryInformationProcess", &NtQueryInformationProcess);
            if (!NtQueryInformationProcess) {
                LOG(ERROR) << "Unable to get NtQueryInformationProcess";
                return NULL;
            }

            PROCESS_BASIC_INFORMATION pbi;
            LONG status = NtQueryInformationProcess(
                handle, ProcessBasicInformation,
                &pbi, sizeof(PROCESS_BASIC_INFORMATION), NULL);
            if (!NT_SUCCESS(status)) {
                LOG(ERROR) << "NtQueryInformationProcess failed";
                return NULL;
            }

            return ::OpenProcess(PROCESS_ALL_ACCESS, TRUE,
                pbi.InheritedFromUniqueProcessId);
        }

        StringType AddQuoteForArg(const StringType& arg)
        {
            // We follow the quoting rules of CommandLineToArgvW.
            // http://msdn.microsoft.com/en-us/library/17w5ykft.aspx
            std::wstring quotable_chars(L" \\\"");
            if (arg.find_first_of(quotable_chars) == std::wstring::npos) {
                // No quoting necessary.
                return arg;
            }

            std::wstring out;
            out.push_back(L'"');
            for (size_t i = 0; i < arg.size(); ++i) {
                if (arg[i] == '\\') {
                    // Find the extent of this run of backslashes.
                    size_t start = i, end = start + 1;
                    for (; end < arg.size() && arg[end] == '\\'; ++end) { }
                    size_t backslash_count = end - start;

                    // Backslashes are escapes only if the run is followed by a double quote.
                    // Since we also will end the string with a double quote, we escape for
                    // either a double quote or the end of the string.
                    if (end == arg.size() || arg[end] == '"') {
                        // To quote, we need to output 2x as many backslashes.
                        backslash_count *= 2;
                    }
                    for (size_t j = 0; j < backslash_count; ++j)
                        out.push_back('\\');

                    // Advance i to one before the end to balance i++ in loop.
                    i = end - 1;
                } else if (arg[i] == '"') {
                    out.push_back('\\');
                    out.push_back('"');
                } else {
                    out.push_back(arg[i]);
                }
            }
            out.push_back('"');

            return out;
        }

    } // namespace

    StringType GetWaitEventName(base::ProcessId pid)
    {
        return base::StringPrintf(L"%s-%d", kWaitEventName, static_cast<int>(pid));
    }

    StringType ArgvToCommandLineString(const StringVector& argv)
    {
        StringType command_line;
        for (const StringType& arg : argv) {
            if (!command_line.empty())
                command_line += L' ';
            command_line += AddQuoteForArg(arg);
        }
        return command_line;
    }

    void RelauncherSynchronizeWithParent()
    {
        base::Process process = base::Process::Current();
        base::win::ScopedHandle parent_process(
            GetParentProcessHandle(process.Handle()));

        // Notify the parent process that it can quit now.
        StringType name = internal::GetWaitEventName(process.Pid());
        base::win::ScopedHandle wait_event(
            ::CreateEventW(NULL, TRUE, FALSE, name.c_str()));
        ::SetEvent(wait_event.Get());

        // Wait for parent process to quit.
        WaitForSingleObject(parent_process.Get(), INFINITE);
    }

    int LaunchProgram(const StringVector& relauncher_args,
        const StringVector& argv)
    {
        base::LaunchOptions options;
        base::Process process = base::LaunchProcess(ArgvToCommandLineString(argv), options);
        return process.IsValid() ? 0 : 1;
    }

} // namespace internal

} // namespace relauncher
